#include "mloop.h"

#define C_R 1.87
#define N_R 1.65
#define O_R 1.40
#define S_R 1.85

float judge_radius(struct atomname atom)
{
	float radius;

	if (atom.atomname[0] == 'C')
		radius = C_R;
	else if (atom.atomname[0] == 'N')
		radius = N_R;
	else if (atom.atomname[0] == 'O')
		radius = O_R;
	else if (atom.atomname[0] == 'S')
		radius = S_R;

	return radius;
}

unsigned char judge_status(struct atomname atom)
{
	unsigned char status;

	status = '\0';
	if (atom.atomname[0] == 'C')
	{
		if (!strncmp (atom.atomname, "C  ", 3))
			;
		else if (!strncmp (atom.resname, "ARG", 3) && !strncmp (atom.atomname, "CZ ", 3))
			;
		else if (!strncmp (atom.resname, "ASN", 3) && !strncmp (atom.atomname, "CG ", 3))
			;
		else if (!strncmp (atom.resname, "ASP", 3) && !strncmp (atom.atomname, "CG ", 3))
			;
		else if (!strncmp (atom.resname, "GLN", 3) && !strncmp (atom.atomname, "CD ", 3))
			;
		else if (!strncmp (atom.resname, "GLU", 3) && !strncmp (atom.atomname, "CD ", 3))
			;
		else if (!strncmp (atom.resname, "HIS", 3) && !strncmp (atom.atomname, "CG ", 3))
			;
		else if (!strncmp (atom.resname, "PHE", 3) && strncmp (atom.atomname, "CA ", 3) && strncmp (atom.atomname, "CB ", 3))
			;
		else if (!strncmp (atom.resname, "TYR", 3) && strncmp (atom.atomname, "CA ", 3) && strncmp (atom.atomname, "CB ", 3))
			;
		else if (!strncmp (atom.resname, "TRP", 3) && strncmp (atom.atomname, "CA ", 3) && strncmp (atom.atomname, "CB ", 3))
			;
		else
			status |= IS_DONO;
	}
	else if (atom.atomname[0] == 'N')
	{
		if (strncmp (atom.resname, "PRO", 3))
		{
			status |= IS_DONO;
			if (!strncmp (atom.resname, "HIS", 3) && strncmp (atom.atomname, "N  ", 3))
				status |= IS_ACCE;
		}
	}
	else if (atom.atomname[0] == 'O')
	{
		status |= IS_ACCE;
		if (!strncmp (atom.resname, "TYR", 3) && !strncmp (atom.atomname, "OH ", 3))
			status |= IS_DONO;
		else if (!strncmp (atom.resname, "SER", 3) && !strncmp (atom.atomname, "OG ", 3))
			status |= IS_DONO;
		else if (!strncmp (atom.resname, "THR", 3) && !strncmp (atom.atomname, "OG1", 3))
			status |= IS_DONO;
	}
	else if (atom.atomname[0] == 'O')
	{
		status |= IS_ACCE;
		if (!strncmp (atom.resname, "CYS", 3))
			status |= IS_DONO;
	}

	if (strncmp (atom.atomname, "CA ", 3) == 0)
	{
		status |= IS_CA;
		status |= IS_BACKBONE;
		status |= IS_MAINCHAIN_AND_CB;
		if (strncmp (atom.resname, "GLY", 3) == 0)
		{
			status |= IS_CB;
		}
	}
	else if (strncmp (atom.atomname, "N  ", 3) == 0 || strncmp (atom.atomname, "C  ", 3) == 0)
	{
		status |= IS_BACKBONE;
		status |= IS_MAINCHAIN_AND_CB;
	}
	else if (strncmp (atom.atomname, "O  ", 3) == 0)
	{
		status |= IS_MAINCHAIN_AND_CB;
	}
	else if (strncmp (atom.atomname, "CB ", 3) == 0)
	{
		status |= IS_MAINCHAIN_AND_CB;
		status |= IS_CB;
	}

	return status;
}

void matrix_inverse(float a[3*3], float result[3*3])
{
	float tmp;

	tmp = (a[0] * a[4] - a[1] * a[3]) * a[8] + (a[2] * a[3] - a[0] * a[5]) * a[7] + (a[1] * a[5] - a[2] * a[4]) * a[6];
	result[0] = (a[4] * a[8] - a[5] * a[7]) / tmp;
	result[3] = (a[5] * a[6] - a[3] * a[8]) / tmp;
	result[6] = (a[3] * a[7] - a[4] * a[6]) / tmp;

	result[1] = (a[2] * a[7] - a[1] * a[8]) / tmp;
	result[4] = (a[0] * a[8] - a[2] * a[6]) / tmp;
	result[7] = (a[1] * a[6] - a[0] * a[7]) / tmp;

	result[2] = (a[1] * a[5] - a[2] * a[4]) / tmp;
	result[5] = (a[2] * a[3] - a[0] * a[5]) / tmp;
	result[8] = (a[0] * a[4] - a[1] * a[3]) / tmp;
}

void linear_solve(float a_data[3*3], float b_data[3], float result[3])
{
	float c_data[3*3];

	matrix_inverse (a_data, c_data);
	result[0] = b_data[0] * c_data[0] + b_data[1] * c_data[1] + b_data[2] * c_data[2];
	result[1] = b_data[0] * c_data[3] + b_data[1] * c_data[4] + b_data[2] * c_data[5];
	result[2] = b_data[0] * c_data[6] + b_data[1] * c_data[7] + b_data[2] * c_data[8];
}

void rotate_matrix(float a[3], float b[3], float dihedral, float a_data[3*3])
{
	float rot_axis[3], r2, f1, f2;

	rot_axis[0] = b[0] - a[0];
	rot_axis[1] = b[1] - a[1];
	rot_axis[2] = b[2] - a[2];
	r2 = rot_axis[0] * rot_axis[0] + rot_axis[1] * rot_axis[1] + rot_axis[2] * rot_axis[2];
	r2 = sqrtf (r2);
	rot_axis[0] /= r2;
	rot_axis[1] /= r2;
	rot_axis[2] /= r2;

	f1 = sinf (dihedral / RADIAN);
	f2 = cosf (dihedral / RADIAN);

	a_data[0] = rot_axis[0] * rot_axis[0] + (1 - rot_axis[0] * rot_axis[0]) * f2;
	a_data[1] = rot_axis[0] * rot_axis[1] * (1 - f2) + rot_axis[2] * f1;
	a_data[2] = rot_axis[0] * rot_axis[2] * (1 - f2) - rot_axis[1] * f1;
	a_data[3] = rot_axis[0] * rot_axis[1] * (1 - f2) - rot_axis[2] * f1;
	a_data[4] = rot_axis[1] * rot_axis[1] + (1 - rot_axis[1] * rot_axis[1]) * f2;
	a_data[5] = rot_axis[1] * rot_axis[2] * (1 - f2) + rot_axis[0] * f1;
	a_data[6] = rot_axis[0] * rot_axis[2] * (1 - f2) + rot_axis[1] * f1;
	a_data[7] = rot_axis[1] * rot_axis[2] * (1 - f2) - rot_axis[0] * f1;
	a_data[8] = rot_axis[2] * rot_axis[2] + (1 - rot_axis[2] * rot_axis[2]) * f2;
}

void rotate_atom(float a[3], float b[3], float c[3], float dihedral)
{
	float a_data[3*3], b_data[3], c_data[3];

	rotate_matrix (a, b, dihedral, a_data);
	b_data[0] = c[0] - a[0];
	b_data[1] = c[1] - a[1];
	b_data[2] = c[2] - a[2];

	c_data[0] = b_data[0] * a_data[0] + b_data[1] * a_data[3] + b_data[2] * a_data[6];
	c_data[1] = b_data[0] * a_data[1] + b_data[1] * a_data[4] + b_data[2] * a_data[7];
	c_data[2] = b_data[0] * a_data[2] + b_data[1] * a_data[5] + b_data[2] * a_data[8];

	c[0] = c_data[0] + a[0];
	c[1] = c_data[1] + a[1];
	c[2] = c_data[2] + a[2];
}

void rotate_atoms(float a[3], float b[3], struct candidate *p, int segid, int flag[LOOPNUM], float dihedral)
{
	float a_data[3*3], b_data[3], c_data[3];
	int i;

	rotate_matrix (a, b, dihedral, a_data);

	for (i = loop.start[loop.seg_start[segid]]; i <= loop.end[loop.seg_end[segid]]; i++)
	{
		if (!flag[i])
			continue;
		b_data[0] = p->coor[i][0] - a[0];
		b_data[1] = p->coor[i][1] - a[1];
		b_data[2] = p->coor[i][2] - a[2];

		c_data[0] = b_data[0] * a_data[0] + b_data[1] * a_data[3] + b_data[2] * a_data[6];
		c_data[1] = b_data[0] * a_data[1] + b_data[1] * a_data[4] + b_data[2] * a_data[7];
		c_data[2] = b_data[0] * a_data[2] + b_data[1] * a_data[5] + b_data[2] * a_data[8];

		p->coor[i][0] = c_data[0] + a[0];
		p->coor[i][1] = c_data[1] + a[1];
		p->coor[i][2] = c_data[2] + a[2];
	}

	b_data[0] = p->dummy_anchor[segid].N[0] - a[0];
	b_data[1] = p->dummy_anchor[segid].N[1] - a[1];
	b_data[2] = p->dummy_anchor[segid].N[2] - a[2];

	c_data[0] = b_data[0] * a_data[0] + b_data[1] * a_data[3] + b_data[2] * a_data[6];
	c_data[1] = b_data[0] * a_data[1] + b_data[1] * a_data[4] + b_data[2] * a_data[7];
	c_data[2] = b_data[0] * a_data[2] + b_data[1] * a_data[5] + b_data[2] * a_data[8];

	p->dummy_anchor[segid].N[0] = c_data[0] + a[0];
	p->dummy_anchor[segid].N[1] = c_data[1] + a[1];
	p->dummy_anchor[segid].N[2] = c_data[2] + a[2];

	b_data[0] = p->dummy_anchor[segid].CA[0] - a[0];
	b_data[1] = p->dummy_anchor[segid].CA[1] - a[1];
	b_data[2] = p->dummy_anchor[segid].CA[2] - a[2];

	c_data[0] = b_data[0] * a_data[0] + b_data[1] * a_data[3] + b_data[2] * a_data[6];
	c_data[1] = b_data[0] * a_data[1] + b_data[1] * a_data[4] + b_data[2] * a_data[7];
	c_data[2] = b_data[0] * a_data[2] + b_data[1] * a_data[5] + b_data[2] * a_data[8];

	p->dummy_anchor[segid].CA[0] = c_data[0] + a[0];
	p->dummy_anchor[segid].CA[1] = c_data[1] + a[1];
	p->dummy_anchor[segid].CA[2] = c_data[2] + a[2];
}

/* type: 1 phi 2 psi */
void set_flag (int segid, int resno, int type, int rot_flag[LOOPNUM])
{
	int i;

	for (i = loop.start[loop.seg_start[segid]] ; i <= loop.start[resno] + 1; i++)
		rot_flag[i] = 0;
	if (type == 1)
	{
		for (i = loop.start[resno] + 2; i < loop.end[resno] + 1; i++)
			rot_flag[i] = 1;
	}
	else
	{
		for (i = loop.start[resno] + 2; i < loop.end[resno] + 1; i++)
			rot_flag[i] = 0;
		rot_flag[loop.start[resno] + 3] = 1;
	}
	for (i = loop.end[resno] + 1; i <= loop.end[loop.seg_end[segid]]; i++)
		rot_flag[i] = 1;
}

int random_rotate(struct candidate *p, int segid)
{
	int resnum, n, rot_flag[LOOPNUM];
	float angle, phi, psi;

	resnum = loop.seg_end[segid] - loop.seg_start[segid] + 1;
	n = drand48 () * resnum + loop.seg_start[segid];
	rand_phi_psi (&phi, &psi, loop.type[n]);
	set_flag (segid, n, 1, rot_flag);
	angle = phi - p->phi[n];
	rotate_atoms (p->coor[loop.start[n]], p->coor[loop.start[n] + 1], p, segid, rot_flag, angle);
	set_flag (segid, n, 2, rot_flag);
	angle = psi - p->psi[n];
	rotate_atoms (p->coor[loop.start[n] + 1], p->coor[loop.start[n] + 2], p, segid, rot_flag, angle);
	p->phi[n] = phi;
	p->psi[n] = psi;

	return n;
}

int random_rotate2(struct candidate *p, int segid, float range)
{
	int resnum, n, rot_flag[LOOPNUM];
	float angle, phi, psi;

	resnum = loop.seg_end[segid] - loop.seg_start[segid] + 1;
	n = drand48 () * resnum + loop.seg_start[segid];
	if (loop.type[n] == 12)
		rand_phi_psi (&phi, &psi, loop.type[n]);
	else
		rand_phi_psi2 (&phi, &psi, p->phi[n], p->psi[n], loop.type[n], range);
	set_flag (segid, n, 1, rot_flag);
	angle = phi - p->phi[n];
	rotate_atoms (p->coor[loop.start[n]], p->coor[loop.start[n] + 1], p, segid, rot_flag, angle);
	set_flag (segid, n, 2, rot_flag);
	angle = psi - p->psi[n];
	rotate_atoms (p->coor[loop.start[n] + 1], p->coor[loop.start[n] + 2], p, segid, rot_flag, angle);
	p->phi[n] = phi;
	p->psi[n] = psi;

	return n;
}

void reinit(struct candidate *p)
{
	int i, j, k, flag, rot_flag[LOOPNUM];
	float angle, phi, psi;

	printf ("Reinitialize...\n");

	for (i = 0; i < 1000; i++)
	{
		for (j = 0; j < loop.nseg; j++)
		{
			for (k = loop.seg_start[j]; k <= loop.seg_end[j]; k++)
			{
				rand_phi_psi (&phi, &psi, loop.type[k]);
				set_flag (j, k, 1, rot_flag);
				angle = phi - p->phi[k];
				rotate_atoms (p->coor[loop.start[k]], p->coor[loop.start[k] + 1], p, j, rot_flag, angle);
				set_flag (j, k, 2, rot_flag);
				angle = psi - p->psi[k];
				rotate_atoms (p->coor[loop.start[k] + 1], p->coor[loop.start[k] + 2], p, j, rot_flag, angle);
				p->phi[k] = phi;
				p->psi[k] = psi;
			}
		}

		flag = 0;
		for (j = 0; j < loop.nseg; j++)
		{
			if (loop_closure2 (p, -1, 100, j) == -1)
			{
				flag = 1;
				break;
			}
		}
		if (flag == 1 || calc_energy1 (p) < 0)
		{
			flag = 1;
			continue;
		}
		else
			break;
	}

	if (flag == 1)
	{
		if (nsaves == 0)
		{
			printf ("Failed in loop closure.\n");
			exit (-1);
		}
		else
		{
			int select;

			select = nsaves * drand48 ();
			memcpy (p, &candidate[select], sizeof (struct candidate));
		}
	}
}

void purturb_mainchain ()
{
	int i, n, flag, segid;

	flag = 0;
	for (i = 0; i < 1000; i++)
	{
		memcpy (&current[1], &current[0], sizeof (struct candidate));
		segid = drand48 () * loop.nseg;
		n = random_rotate (&current[1], segid);
		if (loop_closure2 (&current[1], n, loop.circle, segid) == -1)
			continue;
		if (calc_energy1 (&current[1]) < 0)
			continue;
		flag = 1;
		break;
	}

	if (flag == 1)
		return;

	if (nsaves == 0)
	{
		reinit (&current[1]);
	}
	else
	{
		n = nsaves * drand48 ();
		memcpy (&current[1], &candidate[n], sizeof (struct candidate));
	}
}

void purturb_sidechain ()
{
	int res_no, rotamer_no;
	float tmp1, tmp2;

	while (1)
	{
		res_no = drand48 () * loop.resnum2;
		res_no = loop.resno[res_no];
		do
		{
			rotamer_no = drand48 () * sidechain[loop.type[res_no]].rotamer_num;
		}
		while (rotamer_no == current[1].rotamer[res_no]);

		current[1].rotamer[res_no] = rotamer_no;
		paste_sidechain (current[1].coor[loop.start[res_no]], current[1].coor[loop.start[res_no] + 1], current[1].coor[loop.start[res_no] + 2], &current[1], loop.start[res_no] + 4, 0, loop.type[res_no], current[1].rotamer[res_no], 0);

		tmp1 = calc_energy2 (&current[1], res_no);
		if (tmp1 < 0.0) 
		{
			memcpy (&current[1], &current[0], sizeof (struct candidate));
			continue;
		}
		tmp2 = calc_energy2 (&current[0], res_no);
		current[1].energy = current[1].energy - tmp2 + tmp1;
		current[1].energy2 = current[1].energy2 - tmp2 + tmp1;
		break;
	}
}

float calc_rmsd(struct candidate *p1, struct candidate *p2)
{
	float rmsd;
	int i, j, n, m;

	rmsd = 0.0;
	n = 0;
	for (i = 0; i < loop.nseg; i++)
	{
		for (j = loop.start[loop.seg_start[i]]; j <= loop.end[loop.seg_end[i]]; j++)
		{
			if (loop.loopatom[j].status & IS_MAINCHAIN_AND_CB)
			{
				rmsd += (p1->coor[j][0] - p2->coor[j][0]) * (p1->coor[j][0] - p2->coor[j][0]) + (p1->coor[j][1] - p2->coor[j][1]) * (p1->coor[j][1] - p2->coor[j][1]) + (p1->coor[j][2] - p2->coor[j][2]) * (p1->coor[j][2] - p2->coor[j][2]);
				n ++;
			}
		}
	}
	rmsd /= n;
	rmsd = sqrtf (rmsd);
	return rmsd;
}

void loop_output(char *filename)
{
	FILE *fp;
	int i, n, m, num, flag;
	struct candidate *p;
#ifdef DEBUG
	float rmsd, rmsd_min, energy;

	rmsd_min = 9.9; energy = 9999.9;
#endif

	fp = fopen (filename, "w");
	if (fp == NULL)
	{
		printf ("open %s failed\n", filename);
		exit (-1);
	}
	for (n = 0; n < nsaves; n++)
	{
		p = &(candidate2[n]);
		p->energyavg2 /= repr_num[n];
#ifdef DEBUG
		rmsd = calc_rmsd (&ref, p);
		if (rmsd < rmsd_min)
		{
			energy = p->energy2;
			rmsd_min = rmsd;
		}
		fprintf (fp, "MODEL %4d ENERGY %8.3f %8.3f %8.3f %8.3f E2_AVG %8.3f CLASS_NUM %4d RMSD %8.3f\n", n + 1, p->energy, p->energy1, p->energy2, p->energypmf, p->energyavg2, repr_num[n], rmsd);
#else
		fprintf (fp, "MODEL %4d ENERGY %8.3f %8.3f %8.3f %8.3f E2_AVG %8.3f CLASS_NUM %4d\n", n + 1, p->energy, p->energy1, p->energy2, p->energypmf, p->energyavg2, repr_num[n]);
#endif
		num = 1;
		for (m = 0; m < loop.nseg; m++)
		{
			for (i = loop.start[loop.seg_start[m]]; i <= loop.end[loop.seg_end[m]]; i++)
			{
				fprintf (fp, "ATOM  %5d  %-3s %3s  %4d    %8.3f%8.3f%8.3f  1.00  0.00      PROT\n", num, loop.loopatom[i].name.atomname, loop.loopatom[i].name.resname, loop.loopatom[i].resno, p->coor[i][0], p->coor[i][1], p->coor[i][2]);
				num ++;
			}
		}
		fprintf (fp, "ENDMDL\n");
	}
	fclose (fp);
#ifdef DEBUG
	printf ("Best RMSD: %8.3f %8.3f\t", rmsd_min, energy);
#endif
}
